#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# @Author: renjin@bit.edu.cn
# @Date  : 2024-08-06

"""
【节点名称】：
    COCOResultAssembleNode
【依赖项安装】：
    pip install spirems
【订阅类型】：
    spirecv_msgs::2DTargets （目标检测结果）
【发布类型】：
    无
【构造参数说明】：
    parameter_file (str): 全局参数文件
【节点参数】：
    /work_dir (str): 全局工作空间
    /inputs_fn (str): 全局评估任务输入文件夹
【备注】：
    无
"""

import threading
import os
import json
import argparse
from typing import Union
from queue import Queue
from spirems import Subscriber, Publisher, def_msg, QoS, BaseNode, get_extra_args


class COCOResultAssembleNode(threading.Thread, BaseNode):
    def __init__(
        self,
        job_name: str,
        ip: str = '127.0.0.1',
        port: int = 9094,
        param_dict_or_file: Union[dict, str] = None,
        sms_shutdown: bool = True,
        **kwargs
    ):
        threading.Thread.__init__(self)
        sms_shutdown = True if sms_shutdown in ['True', 'true', '1', True] else False
        BaseNode.__init__(
            self,
            self.__class__.__name__,
            job_name,
            ip=ip,
            port=port,
            param_dict_or_file=param_dict_or_file,
            sms_shutdown=sms_shutdown,
            **kwargs
        )
        print("  \033[32m{}\033[0m = {}".format("sms_shutdown", sms_shutdown))
        self.specified_input_topic = self.get_param("specified_input_topic", "")
        self.launch_next_emit = self.get_param("launch_next_emit", False)
        self.remote_ip = self.get_param("remote_ip", "127.0.0.1")
        self.remote_port = self.get_param("remote_port", 9094)
        self.g_work_dir = self.get_param("/work_dir", "C:/tmp")
        self.g_inputs_fn = self.get_param("/inputs_fn", 'COCOEvalNodeInputs')
        self.params_help()

        self.result_queue = Queue()
        self.queue_pool.append(self.result_queue)

        input_topic_url = '/' + job_name + '/detector/results' if len(self.specified_input_topic) == 0 \
            else self.specified_input_topic
        self._result_reader = Subscriber(
            input_topic_url, 'spirecv_msgs::2DTargets', self.results_callback,
            ip=self.remote_ip, port=self.remote_port, qos=QoS.Reliability
        )

        if self.launch_next_emit:
            self._next_writer = Publisher(
                '/' + job_name + '/launch_next', 'std_msgs::Boolean',
                ip=ip, port=port, qos=QoS.Reliability
            )

        self.avg_time_used = 0
        self.start()

    def release(self):
        BaseNode.release(self)
        self._result_reader.kill()
        if self.launch_next_emit:
            self._next_writer.kill()

    def results_callback(self, msg):
        self.result_queue.put(msg)

    def run(self):
        while self.is_running():
            msg = self.result_queue.get(block=True)
            if msg is None:
                break

            assert 'dataset' in msg and len(msg['dataset']) > 0
            assert 'client_id' in msg and len(msg['client_id']) > 0

            path = os.path.join(self.g_work_dir, msg['dataset'], msg['client_id'])
            if not os.path.exists(path):
                os.makedirs(path)
                self.avg_time_used = 0
            if 'time_used' in msg:
                self.avg_time_used += msg['time_used']

            json_fn = os.path.join(str(path), os.path.basename(msg['file_name']) + '.json')
            with open(json_fn, 'w') as f:
                json.dump(msg, f)

            if self.launch_next_emit:
                next_msg = def_msg('std_msgs::Boolean')
                next_msg['data'] = True
                self._next_writer.publish(next_msg)

            if msg['img_id'] + 1 == msg['img_total']:
                print('Send to Evaluator!')
                inputs_dir = os.path.join(self.g_work_dir, self.g_inputs_fn)
                if not os.path.exists(inputs_dir):
                    os.makedirs(inputs_dir)
                coco_eval_input = dict()
                if self.avg_time_used > 0:
                    coco_eval_input['fps'] = 1.0 / (self.avg_time_used / msg['img_total'])
                else:
                    coco_eval_input['fps'] = 0.0
                coco_eval_input['dataset'] = msg['dataset']
                coco_eval_input['coco_dt_dir'] = path
                coco_eval_input['client_id'] = msg['client_id']
                with open(os.path.join(str(inputs_dir), msg['client_id'] + '.json'), 'w') as f2:
                    json.dump(coco_eval_input, f2)

        self.release()
        print('{} quit!'.format(self.__class__.__name__))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--config',
        type=str,
        default='default_params.json',
        help='SpireCV2 Config (.json)')
    parser.add_argument(
        '--job-name',
        type=str,
        default='live',
        help='SpireCV Job Name')
    parser.add_argument(
        '--ip',
        type=str,
        default='127.0.0.1',
        help='SpireMS Core IP')
    parser.add_argument(
        '--port',
        type=int,
        default=9094,
        help='SpireMS Core Port')
    # args = parser.parse_args()
    args, unknown_args = parser.parse_known_args()
    if not os.path.isabs(args.config):
        current_path = os.path.abspath(__file__)
        params_dir = os.path.join(current_path[:current_path.find('spirecv-pro') + 11], 'params', 'spirecv2')
        args.config = os.path.join(params_dir, args.config)
    print("--config:", args.config)
    print("--job-name:", args.job_name)
    extra = get_extra_args(unknown_args)

    node = COCOResultAssembleNode(args.job_name, param_dict_or_file=args.config, ip=args.ip, port=args.port, **extra)
    node.join()
